import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
import streamlit as st
from PIL import Image
import io


def detect_lanes(uploaded_file):
    # 将上传的文件转换为opencv格式
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    color_img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    gray_img = cv2.imdecode(file_bytes, cv2.IMREAD_GRAYSCALE)

    # 形态学闭运算
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    gray_img = cv2.morphologyEx(gray_img, cv2.MORPH_CLOSE, kernel)

    # 二值化
    _, binary = cv2.threshold(gray_img, 125, 255, cv2.THRESH_BINARY)

    # Canny边缘检测
    canny = cv2.Canny(binary, 50, 125, 3)

    # 霍夫直线检测
    lines = cv2.HoughLinesP(
        canny,
        rho=1,
        theta=np.pi / 180,
        threshold=100,
        minLineLength=10,
        maxLineGap=50
    )

    # 绘制直线
    result = color_img.copy()
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            # 计算角度
            dx = x2 - x1
            dy = y2 - y1
            angle = np.arctan2(float(dy), dx) * 180 / np.pi
            cv2.line(result, (x1, y1), (x2, y2), (255, 0, 0), 1)

    return gray_img, binary, canny, result


def main():
    st.title("车道线检测系统")

    # 文件上传
    uploaded_file = st.file_uploader("选择一张图片", type=['jpg', 'jpeg', 'png'])

    if uploaded_file is not None:
        # 显示原图
        image = Image.open(uploaded_file)
        st.image(image, caption='上传的图片', use_column_width=True)

        # 处理图片
        uploaded_file.seek(0)  # 重置文件指针
        gray_img, binary, canny, result = detect_lanes(uploaded_file)

        # 创建两行展示处理结果
        col1, col2 = st.columns(2)

        with col1:
            st.subheader("灰度图像")
            st.image(gray_img, use_column_width=True)

            st.subheader("二值化图像")
            st.image(binary, use_column_width=True)

        with col2:
            st.subheader("Canny边缘检测")
            st.image(canny, use_column_width=True)

            st.subheader("检测结果")
            st.image(cv2.cvtColor(result, cv2.COLOR_BGR2RGB), use_column_width=True)

        # 添加下载按钮
        result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(result_rgb)
        buf = io.BytesIO()
        pil_image.save(buf, format="JPEG")
        byte_im = buf.getvalue()

        st.download_button(
            label="下载处理结果",
            data=byte_im,
            file_name="result.jpg",
            mime="image/jpeg",
        )


if __name__ == "__main__":
    main()
